from flask import Flask, request, jsonify
import torch
import torchvision.transforms as transforms
from PIL import Image
import io
# 虽然代码中没有显式使用 Net，但它是 torch.load() 能正确重建模型的必要条件
from model import Net  # 关键修改：从原model.py导入Net类

app = Flask(__name__)

# 加载模型（现在使用统一的Net定义）
model = torch.load("net_0.pth", map_location='cpu', weights_only=False)
model.eval()

# 类别映射
class_list = ["飞机", "汽车", "鸟", "猫", "鹿", "狗", "青蛙", "马", "船", "卡车"]

# 图像预处理
transform = transforms.Compose([
    transforms.Resize((32, 32)),
    transforms.ToTensor()
])


class Result:
    def __init__(self, code=200, message='success', data=None):
        self.code = code
        self.message = message
        self.data = data

    def to_dict(self):
        return {
            'code': self.code,
            'message': self.message,
            'data': self.data
        }


@app.route('/cifar10/predict', methods=['POST'])
def predict():
    if 'file' not in request.files:
        return jsonify(Result(400, '请上传图片文件').to_dict())

    file = request.files['file']
    if file.filename == '':
        return jsonify(Result(400, '空文件名').to_dict())

    try:
        # 检查文件类型
        if not file.filename.lower().endswith(('.png', '.jpg', '.jpeg')):
            return jsonify(Result(400, '仅支持PNG/JPG/JPEG格式').to_dict())

        # 读取图像
        image_bytes = file.read()
        image = Image.open(io.BytesIO(image_bytes)).convert('RGB')

        # 预处理
        image_tensor = transform(image).unsqueeze(0)

        # 预测
        with torch.no_grad():
            output = model(image_tensor)
            probabilities = torch.softmax(output, dim=1)[0]
            pred_class = torch.argmax(output).item()
            confidence = probabilities[pred_class].item()

        # 构建响应数据
        result_data = {
            'prediction': class_list[pred_class],
            'confidence': float(confidence),
            'probabilities': {
                cls: float(prob)
                for cls, prob in zip(class_list, probabilities)
            }
        }

        return jsonify(Result(data=result_data).to_dict())

    except Exception as e:
        return jsonify(Result(500, f'预测失败: {str(e)}').to_dict())


if __name__ == '__main__':
    app.run(port=5000)
